using System.Runtime.InteropServices;

namespace Pinwheel.Kernel
{
    public static class ProcessAffinity
    {
        // Windows API
        [DllImport("kernel32.dll")]
        private static extern IntPtr GetCurrentThread();
        
        [DllImport("kernel32.dll", SetLastError = true)]
        private static extern IntPtr SetThreadAffinityMask(IntPtr hThread, IntPtr dwThreadAffinityMask);
        
        private static void SetAffinityWindows(int coreId)
        {
            IntPtr threadHandle = GetCurrentThread();
            IntPtr affinityMask = (IntPtr)(1 << coreId);
            if (SetThreadAffinityMask(threadHandle, affinityMask) == IntPtr.Zero)
            {
                throw new System.ComponentModel.Win32Exception();
            }
        }

        
        // Linux API
        [DllImport("libc", SetLastError = true)]
        private static extern int pthread_setaffinity_np(IntPtr thread, IntPtr cpusetsize, ref CpuSet cpuset);
        
        [DllImport("libc", SetLastError = true)]
        private static extern IntPtr pthread_self();
        
        [StructLayout(LayoutKind.Sequential)]
        struct CpuSet
        {
            private ulong[] _bits;
            public void Zero() { _bits = new ulong[16]; }
            public void Set(int cpu) { _bits[cpu / 64] |= (1UL << (cpu % 64)); }
        }
        
        private static void SetAffinityLinux(int coreId)
        {
            IntPtr thread = pthread_self();
            var cpuSet = new CpuSet();
            cpuSet.Zero();
            cpuSet.Set(coreId);
            int result = pthread_setaffinity_np(thread, (IntPtr)Marshal.SizeOf(cpuSet), ref cpuSet);
            if (result != 0)
            {
                throw new System.ComponentModel.Win32Exception();
            }
        }
        
        
        // MacOSX API
        [DllImport("libc")]
        private static extern uint pthread_mach_thread_np(IntPtr pthread);
        
        [DllImport("libSystem")]
        private static extern int thread_policy_set(uint thread, int flavor, IntPtr policy_info, uint count);
        
        // Mach 内核 API 相关结构和函数
        private const int THREAD_AFFINITY_POLICY = 4;
        
        [StructLayout(LayoutKind.Sequential)]
        struct thread_affinity_policy_data_t
        {
            public int affinity_tag;
        }
        
        // 设置线程亲和性
        static void SetAffinityMacOSX(int core)
        {
            thread_affinity_policy_data_t policy = new thread_affinity_policy_data_t
            {
                affinity_tag = core
            };

            uint machThread = pthread_mach_thread_np(pthread_self());
            IntPtr policyPtr = Marshal.AllocHGlobal(Marshal.SizeOf(policy));
            Marshal.StructureToPtr(policy, policyPtr, false);
            
            thread_policy_set(machThread, THREAD_AFFINITY_POLICY, policyPtr, 1);

            Marshal.FreeHGlobal(policyPtr);
        }
        
        public static void SetAffinity(int coreId)
        {
            if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
            {
                SetAffinityWindows(coreId);
            }
            else if (RuntimeInformation.IsOSPlatform(OSPlatform.Linux))
            {
                SetAffinityLinux(coreId);
            }
            else if (RuntimeInformation.IsOSPlatform(OSPlatform.OSX))
            {
                SetAffinityMacOSX(coreId); 
            }
            else
            {
                throw new PlatformNotSupportedException(); 
            }
        }
    }
}